import win32evtlog
import win32evtlogutil
import datetime
import re
import collections
import sys
import xml.etree.ElementTree as ET
import os
import csv
import logging # 引入日志模块

# 配置日志，用于在控制台输出调试信息
logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')

def filetime_to_datetime(ft):
    """
    将Windows FILETIME（自1601-01-01 UTC以来100纳秒间隔的数量）
    转换为本地时间的Python datetime对象。
    """
    # FILETIME是一个64位值，表示自1601年1月1日UTC以来100纳秒间隔的数量。
    # 要转换为Unix纪元（1970年1月1日UTC），需要减去1601-01-01和1970-01-01之间100纳秒间隔的数量。
    # 这个常数是116444736000000000（秒 * 10,000,000 以获得100纳秒间隔）
    unix_epoch_diff_100ns = 116444736000000000
    
    # 计算自Unix纪元以来的微秒数
    microseconds_since_unix_epoch = (ft - unix_epoch_diff_100ns) // 10
    
    # 从Unix纪元在UTC中创建一个朴素的datetime对象
    dt_utc = datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc) + datetime.timedelta(microseconds=microseconds_since_unix_epoch)
    
    # 转换为本地时区
    return dt_utc.astimezone(datetime.datetime.now().astimezone().tzinfo)

def parse_event_xml(event_xml):
    """
    解析事件的XML数据以提取相关信息。
    """
    try:
        root = ET.fromstring(event_xml)
    except ET.ParseError as e:
        logging.warning(f"无法解析事件XML: {e}. XML内容可能损坏或不完整。")
        return None, None, None
    
    # 定义命名空间，因为EventData下的Data标签有默认命名空间
    namespaces = {'we': 'http://schemas.microsoft.com/win/2004/08/events/event'}
    
    device_instance_id = None
    event_id = None
    time_created = None

    # 提取Event ID
    event_id_element = root.find(".//{*}EventID")
    if event_id_element is not None:
        try:
            event_id = int(event_id_element.text)
        except ValueError:
            logging.warning(f"无法解析EventID: {event_id_element.text}")

    # 提取时间戳
    time_created_element = root.find(".//{*}TimeCreated")
    if time_created_element is not None:
        time_created_str = time_created_element.get("SystemTime")
        try:
            # 移除 'Z' 表示UTC
            time_created = datetime.datetime.fromisoformat(time_created_str.replace('Z', '+00:00'))
            # 转换为本地时区
            time_created = time_created.astimezone(datetime.datetime.now().astimezone().tzinfo)
        except ValueError as e:
            logging.warning(f"无法解析时间戳 '{time_created_str}': {e}")
            time_created = None # 解析失败

    # 提取Device Instance ID
    for data_element in root.findall(".//{*}EventData/{*}Data", namespaces):
        if data_element.get("Name") == "DeviceInstanceId":
            device_instance_id = data_element.text
            break
            
    return event_id, time_created, device_instance_id

def get_usb_info_from_event_logs():
    """
    从Windows事件日志中提取USB设备信息。
    主要关注Microsoft-Windows-Kernel-PnP/Configuration通道中的连接/断开事件。
    """
    usb_devices = collections.defaultdict(lambda: {
        "vid": "N/A",
        "pid": "N/A",
        "serial_number": "N/A",
        "device_type": "从事件日志",
        "description": "N/A (来自事件日志)",
        "friendly_name": "N/A (来自事件日志)",
        "manufacturer": "N/A (来自事件日志)",
        "first_seen_event_log": None,
        "last_disconnected_event_log": None
    })

    # 尝试不同的日志通道名称，以增加兼容性
    log_channels = [
        "Microsoft-Windows-Kernel-PnP/Configuration", # 最常见的 PnP 配置日志
        "System" # System日志也可能包含PnP事件，但通常不那么详细
    ]
    
    # 常用且可靠的USB连接/断开事件ID
    # 20001: 设备已连接
    # 20003: 设备已断开
    # 20000: 设备安装开始 (通用)
    # 20002: 设备安装结束 (通用)
    # 21004: 驱动程序已为设备启动 (更通用，但可能包含USB)
    # 21005: 驱动程序已为设备停止
    event_ids_to_monitor = [20001, 20003, 20000, 20002, 21004, 21005]
    
    logging.info(f"将尝试从以下事件日志通道获取USB设备信息: {', '.join(log_channels)}")
    logging.info(f"将监控以下事件ID: {', '.join(map(str, event_ids_to_monitor))}")

    for log_channel in log_channels:
        logging.info(f"正在查询通道: {log_channel}...")
        try:
            # 构建XPath查询以过滤特定事件ID
            # 使用 EventID IN (id1, id2, ...) 形式
            event_id_filters = " or ".join([f"EventID={eid}" for eid in event_ids_to_monitor])
            query_xml = (
                "<QueryList>"
                "  <Query Id='0' Path='{channel}'>"
                "    <Select Path='{channel}'>*[System[({event_id_filters})]]</Select>"
                "  </Query>"
                "</QueryList>"
            ).format(channel=log_channel, event_id_filters=event_id_filters)

            h_event_log = win32evtlog.EvtQuery(log_channel, win32evtlog.EvtQueryChannelPath, query_xml, None)
            
            # 批量读取事件
            events_read_count = 0
            while True:
                # 读取下一批事件，每次读取 100 个
                events = win32evtlog.EvtNext(h_event_log, 100)
                if not events:
                    break # 没有更多事件了
                
                for event in events:
                    events_read_count += 1
                    try:
                        event_xml = win32evtlog.EvtRender(event, win32evtlog.EvtRenderEventXml)
                    except Exception as render_e:
                        logging.warning(f"无法渲染事件XML: {render_e}. 跳过此事件。")
                        continue

                    event_id, timestamp, device_instance_id = parse_event_xml(event_xml)

                    if timestamp and device_instance_id:
                        # 改进 Device Instance ID 的解析，使其更健壮
                        # 尝试匹配 USB\VID_XXXX&PID_YYYY\SERIAL_NUMBER 或 USB\VID_XXXX&PID_YYYY&MI_ZZ\SERIAL_NUMBER 等
                        # 最重要的部分是 VID_XXXX&PID_YYYY 以及最后的实例ID（通常是序列号或包含序列号）
                        match = re.search(r"(?:USB\\|USBSTOR\\)VID_([0-9A-F]{4})&PID_([0-9A-F]{4})[^\\]*\\([^\\]+)", device_instance_id, re.IGNORECASE)
                        
                        if match:
                            vid = match.group(1).upper()
                            pid = match.group(2).upper()
                            # 序列号可能在'&'后面有额外部分（例如 "&MI_00"），我们只取第一部分
                            serial = match.group(3).split('&')[0].upper()
                            
                            if not serial or serial == "N/A": # 确保序列号有效，避免通用字符串
                                continue # 跳过没有有效序列号的条目

                            device_key = (vid, pid, serial)

                            device_info = usb_devices[device_key]
                            device_info["vid"] = vid
                            device_info["pid"] = pid
                            device_info["serial_number"] = serial

                            if event_id in [20001, 20000, 21004]: # 视为连接或安装事件
                                if device_info["first_seen_event_log"] is None or timestamp < device_info["first_seen_event_log"]:
                                    device_info["first_seen_event_log"] = timestamp
                            elif event_id in [20003, 20002, 21005]: # 视为断开或卸载事件
                                if device_info["last_disconnected_event_log"] is None or timestamp > device_info["last_disconnected_event_log"]:
                                    device_info["last_disconnected_event_log"] = timestamp
                        # else:
                        #     logging.debug(f"未能从DeviceInstanceId中解析VID/PID/Serial: {device_instance_id}")

                if not events_read_count % 1000: # 每处理1000个事件打印一次进度
                    logging.info(f"已处理 {events_read_count} 个事件...")

            logging.info(f"通道 '{log_channel}' 已处理完成。总共读取了 {events_read_count} 个事件。")

        except Exception as e:
            logging.error(f"读取事件日志通道 '{log_channel}' 时出错: {e}")
            logging.error("请确保你以管理员身份运行此脚本。")
            continue # 继续尝试下一个通道

    # 过滤掉没有有效连接时间戳的设备，并转换为列表
    final_devices_data = [info for info in usb_devices.values() if info["first_seen_event_log"] is not None]
    
    return final_devices_data

def main():
    logging.info("正在尝试收集USB设备信息。这需要管理员权限才能完全访问。")
    
    # 1. 从事件日志获取信息
    usb_devices_data = get_usb_info_from_event_logs()
    
    logging.info(f"\n从事件日志中找到 {len(usb_devices_data)} 个唯一的USB设备条目。\n")

    output_data = []
    # 按照“最后断开 (事件日志)”的时间进行升序排序。
    # None值将排在最后（通过将None视为True，其他为False来处理）。
    sorted_devices = sorted(
        usb_devices_data,
        key=lambda info: (
            info["last_disconnected_event_log"] is None, # 将None值排在最后
            info["last_disconnected_event_log"]          # 然后按实际的datetime对象排序
        ) if info["last_disconnected_event_log"] else (
            info["first_seen_event_log"] is None, # 如果没有断开时间，则按首次发现时间排序
            info["first_seen_event_log"]
        )
    )

    for info in sorted_devices:
        # 确保所有字段都存在，即使为空，以保证CSV的一致性
        row = {
            "VID": info.get("vid", "N/A"),
            "PID": info.get("pid", "N/A"),
            "Serial Number": info.get("serial_number", "N/A"),
            "Device Type": info.get("device_type", "从事件日志"),
            "Description": info.get("description", "N/A (来自事件日志)"), 
            "Friendly Name": info.get("friendly_name", "N/A (来自事件日志)"),
            "Manufacturer": info.get("manufacturer", "N/A (来自事件日志)"),
            "First Seen (Event Log)": info["first_seen_event_log"].strftime('%Y-%m-%d %H:%M:%S') if info["first_seen_event_log"] else "N/A",
            "Last Disconnected (Event Log)": info["last_disconnected_event_log"].strftime('%Y-%m-%d %H:%M:%S') if info["last_disconnected_event_log"] else "N/A"
        }
        output_data.append(row)
        
        # 打印到控制台
        print(f"供应商ID (VID): {info['vid']}, 产品ID (PID): {info['pid']}, 序列号: {info['serial_number']}")
        print(f"   类型: {info['device_type']}")
        if info['description'] != "N/A (来自事件日志)": print(f"   描述: {info['description']}")
        if info['friendly_name'] != "N/A (来自事件日志)": print(f"   友好名称: {info['friendly_name']}")
        if info.get('manufacturer') != "N/A (来自事件日志)": print(f"   制造商: {info.get('manufacturer')}")
        print(f"   首次发现 (事件日志): {info['first_seen_event_log'].strftime('%Y-%m-%d %H:%M:%S') if info['first_seen_event_log'] else 'N/A'}")
        print(f"   最后断开 (事件日志): {info['last_disconnected_event_log'].strftime('%Y-%m-%d %H:%M:%S') if info['last_disconnected_event_log'] else 'N/A'}")
        print("-" * 30)

    # 可选：保存到CSV
    if output_data:
        csv_file = "usb_device_event_log_history.csv"
        logging.info(f"\n正在将数据保存到 {csv_file}...")
        try:
            # 获取所有字典中所有可能的字段名，以确保CSV标题完整
            fieldnames = sorted(list(set(k for d in output_data for k in d.keys())))
            with open(csv_file, 'w', newline='', encoding='utf-8') as f:
                writer = csv.DictWriter(f, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(output_data)
            logging.info(f"成功保存到 {csv_file}")
        except IOError as e:
            logging.error(f"保存CSV时出错: {e}")

if __name__ == "__main__":
    # 在Windows上可靠地检查管理员权限
    is_admin = False
    if sys.platform == "win32":
        import ctypes # 仅在 Windows 上导入 ctypes
        try:
            is_admin = ctypes.windll.shell32.IsUserAnAdmin() != 0
        except Exception:
            pass # 如果检查失败，则假定不是管理员
    else: # 对于其他操作系统，root通常是0
        try:
            is_admin = (os.geteuid() == 0)
        except AttributeError: # 非Unix-like系统可能没有geteuid
            pass

    if not is_admin:
        logging.warning("警告：此脚本可能需要管理员权限才能访问所有注册表项和系统日志。")
        logging.warning("如果遇到错误或数据缺失，请尝试以管理员身份运行。")

    main()